import os
from flask import Flask, render_template, request, jsonify, send_file
from werkzeug.utils import secure_filename
import subprocess
import threading
import time
import random

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['TRANSLATED_FOLDER'] = 'translated'
app.config['MAX_CONTENT_LENGTH'] = 16 * 1024 * 1024  # 16MB max file size

# 确保上传和翻译后的文件夹存在
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['TRANSLATED_FOLDER'], exist_ok=True)

# 存储翻译进度的字典
translation_progress = {}

def translate_pdf(file_path, output_path, translation_id, model=None):
    try:
        # 更新进度为开始
        translation_progress[translation_id] = {'progress': 0, 'status': 'running'}
        
        # 构建翻译命令
        if model:
            # 使用DeepSeek模型
            command = f'pdf2zh "{file_path}" -s deepseek:{model} -o "{app.config["TRANSLATED_FOLDER"]}"'
        else:
            # 使用默认的Google翻译
            command = f'pdf2zh "{file_path}" -o "{app.config["TRANSLATED_FOLDER"]}"'
        
        # 执行翻译命令
        process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        
        # 模拟更真实的进度
        progress = 0
        while progress < 95:  # 最多到95%，等待实际完成
            if translation_progress[translation_id]['status'] == 'running':
                # 随机增加1-5的进度
                progress += random.uniform(1, 5)
                # 确保不超过95
                progress = min(95, progress)
                translation_progress[translation_id]['progress'] = progress
                # 根据进度调整等待时间，越往后越慢
                sleep_time = 0.1 + (progress / 100) * 0.3
                time.sleep(sleep_time)
        
        stdout, stderr = process.communicate()
        
        if process.returncode == 0:
            # 只有在实际成功时才显示100%
            translation_progress[translation_id] = {'progress': 100, 'status': 'completed'}
        else:
            translation_progress[translation_id] = {'progress': progress, 'status': 'error', 'message': stderr.decode()}
            
    except Exception as e:
        translation_progress[translation_id] = {'progress': 0, 'status': 'error', 'message': str(e)}

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/about')
def about():
    return render_template('about.html')

@app.route('/upload', methods=['POST'])
def upload_file():
    if 'file' not in request.files:
        return jsonify({'error': '没有文件上传'}), 400
    
    file = request.files['file']
    if file.filename == '':
        return jsonify({'error': '没有选择文件'}), 400
    
    if not file.filename.endswith('.pdf'):
        return jsonify({'error': '只支持PDF文件'}), 400
    
    filename = secure_filename(file.filename)
    file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
    file.save(file_path)
    
    # 生成翻译ID
    translation_id = str(int(time.time()))
    
    # 获取翻译模型选择
    model = request.form.get('model')  # 可以是 'deepseek-chat' 或 'deepseek-reasoner'
    
    # 启动翻译线程
    thread = threading.Thread(
        target=translate_pdf,
        args=(file_path, os.path.join(app.config['TRANSLATED_FOLDER'], f'translated_{filename}'), translation_id, model)
    )
    thread.start()
    
    return jsonify({
        'message': '文件上传成功，开始翻译',
        'translation_id': translation_id
    })

@app.route('/progress/<translation_id>')
def get_progress(translation_id):
    if translation_id not in translation_progress:
        return jsonify({'error': '翻译任务不存在'}), 404
    return jsonify(translation_progress[translation_id])

@app.route('/files')
def list_files():
    uploads = []
    translated = []
    
    # 获取上传的文件
    for file in os.listdir(app.config['UPLOAD_FOLDER']):
        if file.endswith('.pdf'):
            uploads.append({
                'name': file,
                'path': os.path.join(app.config['UPLOAD_FOLDER'], file)
            })
    
    # 获取翻译后的文件
    for file in os.listdir(app.config['TRANSLATED_FOLDER']):
        if file.endswith('.pdf'):
            translated.append({
                'name': file,
                'path': os.path.join(app.config['TRANSLATED_FOLDER'], file)
            })
    
    return jsonify({
        'uploads': uploads,
        'translated': translated
    })

@app.route('/download/<path:filename>')
def download_file(filename):
    try:
        # 根据文件名判断是源文件还是翻译后的文件
        if filename.endswith(('-dual.pdf', '-mono.pdf')):
            directory = app.config['TRANSLATED_FOLDER']
        else:
            directory = app.config['UPLOAD_FOLDER']
        return send_file(os.path.join(directory, filename), as_attachment=True)
    except Exception as e:
        return jsonify({'error': str(e)}), 404

@app.route('/preview/<path:filename>')
def preview_file(filename):
    try:
        # 根据文件名判断是源文件还是翻译后的文件
        if filename.endswith(('-dual.pdf', '-mono.pdf')):
            directory = app.config['TRANSLATED_FOLDER']
        else:
            directory = app.config['UPLOAD_FOLDER']
        return send_file(os.path.join(directory, filename))
    except Exception as e:
        return jsonify({'error': str(e)}), 404

@app.route('/preview_page/<path:filename>')
def preview_page(filename):
    return render_template('preview.html')

@app.route('/get_backgrounds')
def get_backgrounds():
    try:
        backgrounds = []
        res_dir = os.path.join(os.path.dirname(__file__), 'res')
        for file in os.listdir(res_dir):
            if file.lower().endswith(('.jpg', '.jpeg', '.png')):
                backgrounds.append(file)
        return jsonify(backgrounds)
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/res/<path:filename>')
def serve_background(filename):
    res_dir = os.path.join(os.path.dirname(__file__), 'res')
    return send_file(os.path.join(res_dir, filename))

if __name__ == '__main__':
    app.run(debug=True)
